//
// Created by syshen on 2021/6/29.
//

#ifndef OCR_CRAFT_H
#define OCR_CRAFT_H

#include "NvCaffeParser.h"
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "NvUffParser.h"
#include <opencv2/opencv.hpp>

#include "utils/cudaWrapper.h"
#include "utils/ioHelper.h"
#include "TextDetector/base_detect.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30)) // 1 GB
#define MAX_BINDING_NUMS   (2)

class DraftDet : public BaseDetect {
public:
    DraftDet(int device_id, bool trt_serialize = false) : BaseDetect(device_id, trt_serialize) {};

    int init(const std::string &model_file);

    std::vector <std::vector<cv::Point2f>> extract(const cv::Mat &inputs);


    ~DraftDet();

private:
    int serialization(const std::string &model_file, const std::string &engine_file);

    bool getImgSize(const cv::Mat &tmp_, cv::Mat &output, float &ratio_x, float &ratio_y);

    bool Input(cv::Mat &src, float *data);

    std::vector <std::vector<cv::Point2f>> Output(const cv::Mat &txt, const cv::Mat &link,const cv::Mat &input, float &ratio_x, float &ratio_y);

    std::vector <std::vector<cv::Point2f>> Output(const cv::Mat &txt, const cv::Mat &link, float &ratio_x, float &ratio_y);

private:

    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};
    std::unique_ptr <nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    void *bindings[MAX_BINDING_NUMS]{0};
    cudaStream_t stream_t;

    //cv::Mat input;

    int device_id = 0;
    int channels = 3;
    int inputH, inputW;
    bool trt_serialize;

};


#endif //OCR_CRAFT_H
